import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import time
import matplotlib.pyplot as plt
from collections import deque
import fitz
import shutil
import os
import re
import glob
import tqdm
import zipfile
try:
    from moviepy import VideoFileClip
except:
    from moviepy.editor import VideoFileClip

# input: cv2 BGR | output: cv2 BGR
def ehist(img):  # 三维图像直方图均衡处理
    BGR = [img[:, :, i] for i in range(3)]
    equ = [cv2.equalizeHist(t) for t in BGR]
    equ = cv2.merge(equ)
    return equ


def merge(img1, img2, alpha=0.5, gamma=0):  #图片叠加
    beta = 1 - alpha
    return cv2.addWeighted(img1, alpha, img2, beta, gamma)
    
    
def text_size(width, height, text, draw):  # 自动适配表情包字体大小
    if width > height:
        font_size = int(height * 0.8)
    else:
        font_size = int(width * 0.8)
    delta = 4
    while True:
        font = ImageFont.truetype('font.ttf', font_size)
        bbox = draw.textbbox((0, 0), text, font=font)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]
        if text_width < width and text_height < height:
            return font, text_width, text_height, font_size
        if text_width > width or text_height > height:
            font_size -= 2
            continue
        if font_size < 1:
            return 'size error'

def add_words(img, text):  # 表情包底部添加文字
    img = Image.fromarray(img).convert("RGB")

    width, height = img.size
    delta_height = int(height * 0.2)  # 
    new_img = Image.new(img.mode, (width, height + delta_height), color='white')
    new_img.paste(img, (0, 0))
    draw = ImageDraw.Draw(new_img)

    fill = (0, 0, 0, 0)
    offset_y = int(height/45)
    font, text_width, text_height, font_size = text_size(width, delta_height, text, draw)
    x = int((width - text_width) / 2)
    y = height + int((delta_height)/2) - text_height
    draw.text((x, y+offset_y), text, font=font, fill=fill)
    
    return np.array(new_img)


# 计算种子点和其领域的像素值之差
def getGrayDiff(gray, current_seed, tmp_seed):
    return abs(int(gray[current_seed[0], current_seed[1]]) - int(gray[tmp_seed[0], tmp_seed[1]]))


# 区域生长算法
def regional_growth(gray, seeds, x, threshold):
    # 八领域
    connects = [(-1, -1), (0, -1), (1, -1), (1, 0), \
                (1, 1), (0, 1), (-1, 1), (-1, 0)]
    org_h, org_w = gray.shape
    # x = 0.3
    gray = cv2.resize(gray, dsize=None, fx=x, fy=x, interpolation=cv2.INTER_AREA)
    seeds = [(int(i[0]*x), int(i[1]*x)) for i in seeds]
    seedMark = np.zeros((gray.shape))
    height, width = gray.shape
    threshold = 10
    seedque = deque()
    label = 255
    seedque.extend(seeds)

    while seedque:
        # 队列具有先进先出的性质。所以要左删
        current_seed = seedque.popleft()
        seedMark[current_seed[0], current_seed[1]] = label
        for i in range(8):
            tmpX = current_seed[0] + connects[i][0]
            tmpY = current_seed[1] + connects[i][1]
            # 处理边界情况
            if tmpX < 0 or tmpY < 0 or tmpX >= height or tmpY >= width:
                continue

            grayDiff = getGrayDiff(gray, current_seed, (tmpX, tmpY))
            if grayDiff < threshold and seedMark[tmpX, tmpY] != label:
                seedque.append((tmpX, tmpY))
                seedMark[tmpX, tmpY] = label
    seedMark = cv2.resize(seedMark, (org_w, org_h), interpolation=cv2.INTER_CUBIC)
    return seedMark


def get_mask(im):
    # import re
    # print(re.findall(r'\'(.*)\'',str(im)))
    # ['background', 'layers', 'composite']
    return im["layers"][0]


def swap(org_im, mask):  # 回调函数，用来实现gr.ImageEditor的掩码作为gr.ImageEditor的输出
    return mask
    # return cv2.addWeighted(org_im, 0.8, mask, 0.2, 0)


def upload(im):
    return im['background'], im['background']


def crop(mask,erode_iter):  # 对掩码进行腐蚀，并返回掩码的box
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (7, 7))  # ksize=7x7,
    # kernel = np.ones((7, 7), np.uint8)
    mask = cv2.erode(mask, kernel, iterations=erode_iter)
    # mask = cv2.medianBlur(mask, 7)
    # mask = cv2.medianBlur(mask, 3)
    nonzero_indexs = [i.tolist() for i in np.nonzero(mask)]
    
    # Handle case where there are no non-zero elements (empty mask)
    if not nonzero_indexs or len(nonzero_indexs) == 0:
        return [0, 0, mask.shape[0], mask.shape[1]]
    
    # Check that we have valid indices for both dimensions
    try:
        box = [min(i) for i in nonzero_indexs]+[max(j) for j in nonzero_indexs]
        return box
    except ValueError as e:
        # If min() fails due to empty iterables, return full mask bounds
        if "min() iterable argument is empty" in str(e):
            return [0, 0, mask.shape[0], mask.shape[1]]
        raise e


def run_region_grow(im, org_im, ehist_flag,canny_flag, erode_iter, x, threshold):
    mask = im['layers'][0]
    mask_gray = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
    seeds = np.transpose(np.nonzero(mask_gray)).tolist()
    if ehist_flag:
        img = cv2.cvtColor(ehist(im['background']), cv2.COLOR_BGR2GRAY)
    else:
        img = cv2.cvtColor(im['background'], cv2.COLOR_BGR2GRAY)
    if canny_flag:
        img = cv2.Canny(img, 100, 300, apertureSize=7, L2gradient=True)
    region_res = np.uint8(regional_growth(img, seeds, x, threshold))

    im_preview = org_im.copy()
    region_res = cv2.medianBlur(region_res, 3) # 滤除区域生长生成的mask椒盐噪声
    box = crop(region_res, erode_iter)
    for i in range(im_preview.shape[-1]):
        im_preview[:,:,i] *= region_res!=0

    im_preview = np.dstack([im_preview, (cv2.medianBlur(region_res*255, 3)!=0)*255])
    # im_preview = region_res*255
    im_preview = im_preview[box[0]:box[2]-2, box[1]:box[3]-2]
    # im_preview = cv2.multiply(img, region_res)
    return region_res, im_preview


def ComputeHist(img):
    h,w = img.shape
    hist, bin_edge = np.histogram(img.reshape(1,w*h), bins=list(range(257)))
    return hist
    
def ComputeMinLevel(hist, rate, pnum):
    sum = 0
    for i in range(256):
        sum += hist[i]
        if (sum >= (pnum * rate * 0.01)):
            return i
            
def ComputeMaxLevel(hist, rate, pnum):
    sum = 0
    for i in range(256):
        sum += hist[255-i]
        if (sum >= (pnum * rate * 0.01)):
            return 255-i
            
def LinearMap(minlevel, maxlevel):
    if (minlevel >= maxlevel):
        return []
    else:
        newmap = np.zeros(256)
        for i in range(256):
            if (i < minlevel):
                newmap[i] = 0
            elif (i > maxlevel):
                newmap[i] = 255
            else:
                newmap[i] = (i-minlevel)/(maxlevel-minlevel) * 255
        return np.array(newmap)
# Automatic Color Equalization Dehazing
def ACED(img):
    h,w,d = img.shape
    newimg = np.zeros([h,w,d])
    for i in range(d):
        # imgmin = np.min(img[:,:,i])
        # imgmax = np.max(img[:,:,i])
        imghist = ComputeHist(img[:,:,i])
        minlevel = ComputeMinLevel(imghist, 8.3, h*w)
        maxlevel = ComputeMaxLevel(imghist, 2.2, h*w)
        newmap = LinearMap(minlevel,maxlevel)
        # print(minlevel, maxlevel)
        if (newmap.size ==0 ):
            continue
        for j in range(h):
            newimg[j,:,i] = newmap[img[j,:, i]]
    return np.array(newimg, dtype='uint8')
# Automatic Color Equalization
def ACE(img, clipLimit, x, y):
    img = np.array(img)
    RGB = [img[:, :, i] for i in range(3)]
    clahe = cv2.createCLAHE(clipLimit=clipLimit, tileGridSize=(x,y))
    equ = [clahe.apply(i) for i in RGB]
    equ = cv2.merge(equ)
    return equ

def pdf2img(pdf_path, x=1.0):
    doc = fitz.open(pdf_path)  # 打开pdf
    pbar = tqdm.tqdm(total=(len(doc)))
    count = 1
    # pdfName = '.'.join(os.path.basename(pdf_path).split('.')[:-1])
    pdfName = '.'.join(pdf_path.split('.')[:-1])
    if not os.path.exists(pdfName):
        os.mkdir(pdfName)
    else:
        if os.name == 'nt':
            os.system(f'del /s /q {pdfName}')
        else:
            os.system(f'rm -rf {pdfName}')
        os.mkdir(pdfName)

    img_dir = pdfName+'/'
    images = []
    savePath = ''
    for page in doc:  # 遍历pdf的每一页
        zoom_x = x  # 设置每页的水平缩放因子
        zoom_y = x  # 设置每页的垂直缩放因子
        mat = fitz.Matrix(zoom_x, zoom_y)
        pix = page.get_pixmap(matrix=mat)
        savePath = img_dir+os.path.basename(pdf_path)[:-4]+f'{count}.jpg'
        pix.save(savePath)     # 保存
        images.append(Image.open(savePath))
        count += 1
        pbar.update(1)
    if len(images) == 1:
        # print('------------')
        # print(savePath)
        # print('pdfName', pdfName)
        # print('pdf_path', pdf_path)
        # print('------------')
        return savePath, images
    files = glob.glob(img_dir+'*')
    zipFile = genNewPath(pdf_path, 'zip')
    with zipfile.ZipFile(zipFile, 'w') as zipf:
        for file in files:
            zipf.write(file, file)
    return zipFile, images

def img2pdf(src):
    with zipfile.ZipFile(src, 'r') as zipf:
        dirlist = []
        for file in zipf.filelist:
            dirlist.append(re.findall(r"filename='(.*)/.*", str(file))[0])
        dirlist = set(dirlist)
        print(dirlist)
        zipf.extractall()
        for i in dirlist:
            tmpdir = i+'/'
            imgs = glob.glob(tmpdir+'*')
            imgs.sort(key=lambda s:int(re.findall(r'(\d+).', ''.join(s.split('.')))[-1]))
            images = []
            for img in imgs:
                image = Image.open(img)
                images.append(image.convert("RGB"))
            timeStr = time.strftime("%Y%m%d%H%M%S")
            images[0].save(f'{i}-{timeStr}.pdf', save_all=True, append_images=images[1:])
            if os.name == 'nt':
                os.system(f'del /s /q {i}')
                os.system(f'rmdir {i}')
            else:
                os.system(f'rm -rf {i}')

'''
def generate_file(file_obj, tmpdir, x=1.0):
    tmpdir = re.findall(r'\'(.*)\'',str(tmpdir))[0]
    print('临时文件夹地址：{}'.format(tmpdir))
    print('上传文件的地址：{}'.format(file_obj.name)) # 输出上传后的文件在gradio中保存的绝对地址

    #获取到上传后的文件的绝对路径后，其余的操作就和平常一致了

    # 将文件复制到临时目录中
    shutil.copy(file_obj.name, tmpdir)

    # 获取上传Gradio的文件名称

    FileName=os.path.basename(file_obj.name)

    # 获取拷贝在临时目录的新的文件地址
    NewfilePath=os.path.join(tmpdir,FileName)
    print(NewfilePath)

    # 打开复制到新路径后的文件
    with open(NewfilePath, 'rb') as file_obj:

        #在本地电脑打开一个新的文件，并且将上传文件内容写入到新文件
        outputPath=os.path.join(tmpdir,"New"+FileName)
        # with open(outputPath,'wb') as w:
        #     w.write(file_obj.read())
    img_path = os.path.join(tmpdir, FileName.split(".pdf")[0]+".png")
    pdf2img(file_obj.name, img_path, x)
    return img_path

'''

def genNewPath(Path, suffix=None):  # 自动给原文件添加时间字符串，从而避免文件名冲突，造成数据覆盖。suffix可修改新文件路径后缀名
    srcName = os.path.basename(Path).split('.')
    if suffix == None:
        suffix = srcName[-1]  # 后缀
    srcName = ''.join(srcName[0:-1])  # 文件名
    srcDir = os.path.dirname(Path)
    timeStr = time.strftime("%Y%m%d%H%M%S")
    count = 0
    index = ''
    while os.path.exists(f'{srcDir}/{srcName}-{timeStr}{index}.{suffix}'):
        index = f'-{count}'
        count += 1
    return f'{srcDir}/{srcName}-{timeStr}{index}.{suffix}'

def getAudio(video):
    my_clip = VideoFileClip(video) #  此处为视频的绝对路径。
    audio = genNewPath(video, 'mp3')
    my_clip.audio.write_audiofile(audio)  # 此处为输出的音频名称
    return audio



